#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sqlite3.h>

#define PORT    2580
#define PORT_2  2581
#define BUFFER_SIZE 1024
#define DATABASE_FILE "user_db.db"

// 函数声明
void handle_client(int client_socket);
int check_user_exists(const char *username);
int register_user(const char *username, const char *password);
int login_user(const char *username, const char *password);
int open_db(sqlite3 **db);
int execute_sql(sqlite3 *db, const char *sql);

int main() {
    pid_t pid;
    pid = fork();
    if(-1 == pid) {
        perror("fork failed");
        return 1;
    }

    int server_socket, client_socket;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);

    sqlite3 *db;
    if (open_db(&db) != 0) {
        fprintf(stderr, "Failed to open the database.\n");
        exit(EXIT_FAILURE);
    }

    // 创建服务器套接字
    if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 设置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);

    if (pid == 0) {
        // 子进程用不同的端口
        server_addr.sin_port = htons(PORT_2);
    }

    // 绑定套接字
    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("Bind failed");
        close(server_socket);
        exit(EXIT_FAILURE);
    }

    // 监听连接
    if (listen(server_socket, 5) == -1) {
        perror("Listen failed");
        close(server_socket);
        exit(EXIT_FAILURE);
    }

    if (pid > 0) {
        printf("Server listening on port %d...\n", PORT);
    }

    if (pid == 0) {
        printf("Server listening on port %d...\n", PORT_2);
    }

    while (1) {
        // 接受客户端连接
        client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &client_addr_len);
        if (client_socket == -1) {
            perror("Accept failed");
            continue;
        }

        printf("Client connected.\n");

        // 处理客户端请求
        handle_client(client_socket);

        // 关闭客户端连接
        close(client_socket);
        printf("Client disconnected.\n");
    }

    sqlite3_close(db);
    close(server_socket);
    return 0;
}

// 打开 SQLite 数据库
int open_db(sqlite3 **db) {
    int rc = sqlite3_open(DATABASE_FILE, db);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "Cannot open database: %s\n", sqlite3_errmsg(*db));
        return -1;
    }

    // 创建用户表（如果尚未创建）
    const char *create_table_sql = "CREATE TABLE IF NOT EXISTS users ("
                                   "id INTEGER PRIMARY KEY AUTOINCREMENT, "
                                   "username TEXT UNIQUE NOT NULL, "
                                   "password TEXT NOT NULL);";
    return execute_sql(*db, create_table_sql);
}

// 执行 SQL 查询或命令
int execute_sql(sqlite3 *db, const char *sql) {
    char *err_msg = NULL;
    int rc = sqlite3_exec(db, sql, 0, 0, &err_msg);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "SQL error: %s\n", err_msg);
        sqlite3_free(err_msg);
        return -1;
    }
    return 0;
}

// 处理客户端请求
void handle_client(int client_socket) {
    char buffer[BUFFER_SIZE];

    // 持续接收客户端请求，直到客户端断开连接
    while (1) {
        ssize_t bytes_read = read(client_socket, buffer, sizeof(buffer) - 1);
        if (bytes_read <= 0) {
            if (bytes_read == 0) {
                // 客户端关闭连接
                printf("Client disconnected.\n");
            } else {
                perror("Read failed");
            }
            break;  // 退出循环，关闭连接
        }

        buffer[bytes_read] = '\0';  // 确保字符串结束
        printf("Received: %s\n", buffer);

        // 解析客户端请求
        char *command = strtok(buffer, ":");
        char *username = strtok(NULL, ":");
        char *password = strtok(NULL, ":");

        if (command == NULL || username == NULL || password == NULL) {
            write(client_socket, "INVALID COMMAND FORMAT", strlen("INVALID COMMAND FORMAT"));
            continue;  // 如果命令格式无效，继续接收下一条请求
        }

        if (strcmp(command, "REGISTER") == 0) {
            if (check_user_exists(username)) {
                write(client_socket, "ALREADY RGS", strlen("ALREADY RGS"));
            } else if (register_user(username, password)) {
                write(client_socket, "SUCCESS", strlen("SUCCESS"));
            } else {
                write(client_socket, "ERROR", strlen("ERROR"));
            }
        } else if (strcmp(command, "LOGIN") == 0) {
            if (login_user(username, password)) {
                write(client_socket, "LOGIN SUCCESS", strlen("LOGIN SUCCESS"));
            } else {
                write(client_socket, "LOGIN FAILED", strlen("LOGIN FAILED"));
            }
        } else {
            write(client_socket, "INVALID COMMAND", strlen("INVALID COMMAND"));
        }
    }

    // 客户端断开连接后，关闭连接
    close(client_socket);
}

// 检查用户是否已存在
int check_user_exists(const char *username) {
    sqlite3 *db;
    if (open_db(&db) != 0) return 0;

    sqlite3_stmt *stmt;
    const char *query = "SELECT COUNT(*) FROM users WHERE username = ?";
    sqlite3_prepare_v2(db, query, -1, &stmt, 0);
    sqlite3_bind_text(stmt, 1, username, -1, SQLITE_STATIC);

    int rc = sqlite3_step(stmt);
    int exists = 0;
    if (rc == SQLITE_ROW) {
        exists = sqlite3_column_int(stmt, 0);
    }

    sqlite3_finalize(stmt);
    sqlite3_close(db);
    return exists > 0;
}

// 注册用户，将用户名和密码存入数据库
int register_user(const char *username, const char *password) {
    sqlite3 *db;
    if (open_db(&db) != 0) return 0;

    sqlite3_stmt *stmt;
    const char *insert_sql = "INSERT INTO users (username, password) VALUES (?, ?)";
    sqlite3_prepare_v2(db, insert_sql, -1, &stmt, 0);
    sqlite3_bind_text(stmt, 1, username, -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, password, -1, SQLITE_STATIC);

    int rc = sqlite3_step(stmt);
    sqlite3_finalize(stmt);
    sqlite3_close(db);

    return rc == SQLITE_DONE;
}

// 登录用户，检查用户名和密码是否匹配
int login_user(const char *username, const char *password) {
    sqlite3 *db;
    if (open_db(&db) != 0) return 0;

    sqlite3_stmt *stmt;
    const char *query = "SELECT password FROM users WHERE username = ?";
    sqlite3_prepare_v2(db, query, -1, &stmt, 0);
    sqlite3_bind_text(stmt, 1, username, -1, SQLITE_STATIC);

    int rc = sqlite3_step(stmt);
    int login_success = 0;
    if (rc == SQLITE_ROW) {
        const char *stored_password = (const char *)sqlite3_column_text(stmt, 0);
        if (strcmp(stored_password, password) == 0) {
            login_success = 1;
        }
    }

    sqlite3_finalize(stmt);
    sqlite3_close(db);
    return login_success;
}

